package com.server.questionbank.snframe.controller;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.server.questionbank.snframe.domain.Ollama.Chat;
import com.server.questionbank.snframe.domain.Ollama.ChatType;
import com.server.questionbank.snframe.domain.Ollama.OllamaRespone;
import com.server.questionbank.snframe.domain.ResDTO.ResponseResult;
import com.server.questionbank.snframe.domain.User.TUser;
import com.server.questionbank.snframe.service.OllamaService.OllamaService;
import com.server.questionbank.snframe.service.UsersService.UsersService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import reactor.core.publisher.Sinks;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/ollama")
public class OllamaController {
    @Autowired
    private OllamaService ollamaService;

    @Autowired
    private UsersService usersService;

    // 创建会话
    @PostMapping("/createChat")
    public ResponseResult<OllamaRespone> createChat(@RequestBody(required = false) String greetings) {
        TUser user = usersService.getLoginUser(); // 判断用户是否登录
        if (user != null) {
            // 用户已登录，创建一个长期有效的会话（在Redis中存储长期，以用户ID为键）
            return ResponseResult.success(ollamaService.createChat(ChatType.LongTerm, greetings, user.getUserId()));
        }
        // 创建一个临时的会话（在Redis中存储30分钟，以最后一次回答开始计算）
         return ResponseResult.success(ollamaService.createChat(ChatType.Temporary.withValue(30), greetings, -1L));
    }

    // 获取会话列表
    @PostMapping("/getChats")
    public ResponseResult<List<Chat>> getChats() {
        TUser user = usersService.getLoginUser(); // 判断用户是否登录
        if (user != null) {
            // 用户已登录，获取长期有效的会话列表
            return ResponseResult.success(ollamaService.getChats(user.getUserId()));
        }
        // 用户未登录，获取临时会话列表
        return ResponseResult.error("未登录，请先登录");
    }

    // 发送消息
    @PostMapping("/sendMessage")
    public ResponseEntity<StreamingResponseBody> sendMessage(@RequestBody OllamaRespone input) {
        // 创建可管理的流
        Sinks.Many<OllamaRespone> sink = Sinks.many().unicast().onBackpressureBuffer();

        // OllamaRespone 中如果没有传入sessionId，则创建一个临时会话。该会话为一次性的对话，即一次对话结束即销毁。
        // 触发 Flux 消息流
        ollamaService.sendGeneralMessage(input).subscribe(sink::tryEmitNext, sink::tryEmitError, sink::tryEmitComplete);

        StreamingResponseBody responseBody = outputStream -> {
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                sink.asFlux().toStream().forEach(response -> {
                    try {
                        String jsonResponse = objectMapper.writeValueAsString(response) + "\n";
                        outputStream.write(jsonResponse.getBytes(StandardCharsets.UTF_8));
                        outputStream.flush();
                    } catch (IOException e) {
                        System.err.println("写入流数据失败: " + e.getMessage());
                    }
                });
            } catch (Exception e) {
                System.err.println("流式输出异常: " + e.getMessage());
                e.printStackTrace();
            } finally {
                try {
                    outputStream.close(); // 关闭流，确保资源释放
                } catch (IOException ignored) {
                    System.err.println("关闭流失败: " + ignored.getMessage());
                }
            }
        };

        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_JSON)
                .header(HttpHeaders.CONNECTION, "keep-alive")  // 维持长连接
                .body(responseBody);
    }

    // 获取总记录数
    @PostMapping("/getTotalRecords")
    public ResponseResult<Long> getTotalRecords(@RequestBody OllamaRespone sessionId) {
        return ollamaService.getTotalRecords(sessionId.getSessionId());
    }

    // 获取聊天记录
    @PostMapping("/getChatHistory")
    public ResponseResult<List<Map<String, String>>> getChatHistory(@RequestBody OllamaRespone sessionId) {
        return ollamaService.getChatHistory(sessionId.getSessionId());
    }

    // 删除会话
    @PostMapping("/deleteChat")
    public ResponseResult<Boolean> deleteChat(@RequestBody OllamaRespone sessionId) {
        Long userId = -1L;
        TUser user = usersService.getLoginUser(); // 判断用户是否登录
        if (user != null) userId = user.getUserId();
        return ollamaService.deleteChat(sessionId.getSessionId(), userId);
    }
}
